import os
import torch
import numpy as np
from PIL import Image
from torch.utils import data
import torchvision.transforms as transforms

imagenet_preprocess = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])


class JIDataset(data.Dataset):
    def __init__(self, split="", preprocessing=None, augmentation=None):
        # train:    4736 = 64 * 74
        # val:      1036 = 4 * 259
        # test:     2416 = 16 * 151
        assert split in ["train", "val", "test"]
        #
        root_dir = "/data/zze/data/JI"
        if not os.path.exists(root_dir):
            root_dir = "/home/hjr/zze/data/JI"
        root_dir = os.path.join(root_dir, split)
        self.images_fps = []
        self.masks_fps = []
        for ne in os.listdir(root_dir + "/image"):
            self.images_fps.append(os.path.join(root_dir + "/image", ne))
            self.masks_fps.append(os.path.join(root_dir + "/label", ne))
        self.augmentation = augmentation
        self.preprocessing = preprocessing

    def __getitem__(self, i):
        image = np.array(Image.open(self.images_fps[i]))
        mask = np.array(Image.open(self.masks_fps[i]))
        mask = (mask != 0).astype('uint8')

        if self.augmentation:
            sample = self.augmentation(image=image, mask=mask)
            image, mask = sample['image'], sample['mask']

        image = imagenet_preprocess(image)
        mask = torch.from_numpy(mask).long()

        return image, mask

    def __len__(self):
        return len(self.images_fps)
